package command

import (
	"context"
	"github.com/spf13/cobra"
	"os"
	"strings"
	"sync"

	"github.com/tikv/client-go/config"
	"github.com/tikv/client-go/key"
	"github.com/tikv/client-go/txnkv"
)

var (
	lock *sync.Mutex = &sync.Mutex{}
	instance *TiKVClient
)

type TiKVClient struct {
	TxnCli *txnkv.Client
}

// KV represents a Key-Value pair.
type KV struct {
	K, V []byte
}

func GetInstance(cmd *cobra.Command) *TiKVClient {
	if instance == nil{
		lock.Lock()
		defer lock.Unlock()

		addrs, err := cmd.Flags().GetString("pd")
		if err != nil {
			cmd.Println("get pd address failed, should set flag with '-u'")
			os.Exit(1)
		}
		eps := strings.Split(addrs, ",")
		TxnCli, err := txnkv.NewClient(context.TODO(), eps, config.Default())
		if err != nil {
			panic(err)
		}
		instance = &TiKVClient{TxnCli}
	}
	return instance
}

func (c *TiKVClient) TxGet(k []byte) (KV, error) {
	//创建事务客户端
	tx, err := c.TxnCli.Begin(context.TODO())
	if err != nil {
		return KV{}, err
	}
	v, err := tx.Get(context.TODO(), k)
	if err != nil {
		return KV{}, err
	}
	return KV{K: k, V:v}, nil
}

func (c *TiKVClient) TxPut(kv []KV) error {
	//创建事务客户端
	tx, err := c.TxnCli.Begin(context.TODO())
	if err != nil {
		return err
	}
	//提交回滚延迟调用
	defer func() {
		if err == nil {
			err = tx.Commit(context.Background())
		}
		if err != nil {
			tx.Rollback()
		}
	}()
	//解析提交数据
	for _, tmpkv := range kv {
		err = tx.Set(tmpkv.K, tmpkv.V)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *TiKVClient) TxDelete(keys [][]byte) error {
	tx, err := c.TxnCli.Begin(context.TODO())
	if err != nil {
		return err
	}
	defer func() {
		if err == nil {
			err = tx.Commit(context.Background())
		}
		if err != nil {
			tx.Rollback()
		}
	}()

	for _, key := range keys {
		err := tx.Delete(key)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *TiKVClient) TxScan(keyPrefix []byte, upperBound []byte, limit int) ([]KV, error) {

	tx, err := c.TxnCli.Begin(context.TODO())
	if err != nil {
		return nil, err
	}

	it, err := tx.Iter(context.TODO(), key.Key(keyPrefix), key.Key(upperBound))
	if err != nil {
		return nil, err
	}
	defer it.Close()
	var ret []KV
	for it.Valid() && limit > 0 {
		ret = append(ret, KV{K: it.Key(), V: it.Value()})
		limit--
		if err := it.Next(context.TODO()); err != nil && it.Valid() {
			return nil, err
		}
	}
	return ret, nil
}

func (c *TiKVClient) Delall(keyPrefix []byte, upperBound []byte) (int64, error) {
	tx, err := c.TxnCli.Begin(context.TODO())
	if err != nil {
		return 0, err
	}

	it, err := tx.Iter(context.TODO(), key.Key(keyPrefix), key.Key(upperBound))
	if err != nil {
		return 0, err
	}
	defer it.Close()

	var count int64
	for it.Valid() {
		tx.Delete(it.Key()[:])
		it.Next(context.TODO())
	}
	err = tx.Commit(context.Background())
	return count, err
}